# 引入必要的包
import os
import numpy as np
from skimage.io import imread
from skimage.transform import resize
from sklearn.model_selection import train_test_split
import yaml

from util import get,dump,preprocess_image

# 读取配置文件中的信息
train_dir = get("train")  # 获取 训练数据路径
char_styles = get("char_styles")  # 获取 字符样式列表，注意: 必须是列标
new_size = tuple(get("new_size"))  # 获取 新图像大小元组, 注意: 必须包含h和w

# 生成X,y
print("# 读取训练数据并进行预处理")
X = []
y = []

# 遍历所有字符风格
for char_style in char_styles:
    # 遍历1到1000的整数
    for i in range(1, 1001):
        # 拼接字符风格和数字，生成图片路径
        img_path = os.path.join(train_dir, f"train_{char_style}_{i:04d}.jpg")
        img_path = os.path.abspath(img_path)
        # 使用预处理函数处理图片
        img = preprocess_image(img_path, new_size)
        # 将处理后的图片添加到X列表中
        X.append(img)
        # 将字符风格在char_styles中的索引值添加到y列表中
        y.append(char_styles.index(char_style))

# 将列表转换为numpy数组
X = np.array(X)
y = np.array(y)

# 分割测试集和训练集
print("# 将数据按 80% 和 20% 的比例分割")
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 打印样本维度和类型信息
print("X_train: ", X_train.shape, X_train.dtype)  # 训练集特征的维度和类型
print("X_test: ", X_test.shape, X_test.dtype)  # 测试集特征的维度和类型
print("y_train: ", y_train.shape, y_train.dtype)  # 训练集标签的维度和类型
print("y_test: ", y_test.shape, y_test.dtype)  # 测试集标签的维度和类型

# 序列化分割后的训练和测试样本
os.makedirs('./Xys', exist_ok=True)
Xy = (X_train, X_test, y_train, y_test)
dump(Xy, "(X_train, X_test, y_train, y_test)", './Xys/Xy')